from selenium import webdriver
import json
import os
import sys
from selenium.webdriver.chrome.options import Options
from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.by import By
options = webdriver.ChromeOptions()

import time
import re
import yaml

class Myclaude():
    def __init__(self):
        print("初始化")
        options.add_argument(f"user-data-dir={os.path.dirname(os.path.abspath(__file__))}/claude")


        self.br = webdriver.Chrome(options=options)
        self.br.maximize_window()
        self.br.get("https://poe.com/")
        self.config_path=f"{os.path.dirname(os.path.abspath(__file__))}/config.yaml"
        self.cookies_path=f"{os.path.dirname(os.path.abspath(__file__))}/claudecookies_test.txt"
        self.config=self.load_config()
        self.refunction=re.compile("[^\\u4e00-\\u9fa5^a-z^A-Z^0-9^\'!\"#$%&\'()*+,\-./:;<=>?@\[\\\]^_`{|}~''＂＃＄％＆＇（）＊＋，－／：；＜＝＞]")
        self.currect_user=None
        self.login()
        pass
    def load_config(self):
        with open(self.config_path, 'rb') as f:
            CONFIG = yaml.load(f.read(), Loader=yaml.FullLoader)
        return CONFIG
    def save_config(self):
        with open(self.config_path, 'w') as f:
            yaml.dump(self.config,f,encoding="utf8",allow_unicode=True)
        pass
    def login(self):
        try:
            if not os.path.exists(self.cookies_path):
                print("请先进行登录，然后按回车键")
                a=input()
                dictCookies = self.br.get_cookies()
                jsonCookies = json.dumps(dictCookies)

                with open(self.cookies_path, 'w') as f:
                    f.write(jsonCookies)

            else:
                with open(self.cookies_path, 'r', encoding='utf8') as f:
                    listCookies = json.loads(f.read())
                for cookie in listCookies:
                    self.br.add_cookie(cookie)
                self.br.refresh()  # 刷新网页,cookies才成功
                time.sleep(0.5)
        except Exception as e:
            print(e)

    def restr(self,txt):

        return self.refunction.sub(" ",txt)

    def meum(self):
        e_menu = self.br.execute_script(
            'return document.querySelectorAll(\'img[alt="Open menu"]\')')
        e_menu[0].click()
        time.sleep(0.5)

    def get_url(self,name:str):
        #todo 根据每个人的微信名，获得相应聊天URL

        try :
            currect_user = self.restr(name)
            users=self.config["users"]

            #判断是否返回主页，重新进入
            if self.currect_user is None:
                self.currect_user=currect_user
            else:
                if currect_user==self.currect_user:
                    return
                else:
                    self.meum()

            #判断是否是历史用户
            if currect_user not in users.keys():#新用户，新建信息
                self.config["users"][currect_user]=""
                self.save_config()

                e_input = self.br.execute_script(
                    'return document.querySelector(\'p[data-placeholder="Message Claude or search past chats..."]\')')

                e_input.send_keys(Keys.CONTROL, "a")
                e_input.send_keys(Keys.BACKSPACE)

                e_input.send_keys(f"i am {currect_user}")
                e_input.send_keys(Keys.ENTER)
                time.sleep(1)
            else:#老用户，转到历史聊天记录
                e_users = self.br.execute_script(
                    'return document.querySelectorAll(\'a[class="absolute inset-0"]\')')
                flag=False
                for i in e_users:
                    if f"i am {currect_user}" == i.text[2:]:
                        i.click()
                        flag=True
                        time.sleep(1)
                        break
                if not flag: #历史聊天已过期，新建聊天
                    e_input = self.br.execute_script(
                        'return document.querySelector(\'p[data-placeholder="Message Claude or search past chats..."]\')')

                    e_input.send_keys(Keys.CONTROL, "a")
                    e_input.send_keys(Keys.BACKSPACE)

                    e_input.send_keys(f"i am {currect_user}")
                    e_input.send_keys(Keys.ENTER)
                    time.sleep(1)
                    pass
        except Exception as e:
            print(e)
            self.reboot()
            self.get_url(name)

        pass
    def wait(self):
        try :
            n = 50
            # 等待回复
            while n:
                e_retry = self.br.execute_script(
                    'return document.querySelectorAll(\'button[class="flex flex-row gap-1 items-center hover:bg-stone-200 p-1 py-0.5 rounded-md transition-opacity delay-100 text-xs"]\')')
                e_refush = self.br.execute_script(
                    'return document.querySelectorAll(\'path[d="M224,128a96,96,0,0,1-94.71,96H128A95.38,95.38,0,0,1,62.1,197.8a8,8,0,0,1,11-11.63A80,80,0,1,0,71.43,71.39a3.07,3.07,0,0,1-.26.25L44.59,96H72a8,8,0,0,1,0,16H24a8,8,0,0,1-8-8V56a8,8,0,0,1,16,0V85.8L60.25,60A96,96,0,0,1,224,128Z"]\')')


                if "Retry" in [i.text for i in e_retry] or len(e_refush)>0:
                    # print("回复好了")
                    break
                else:
                    time.sleep(1)
        except Exception as e:
            print(e)
            time.sleep(20)
        pass
    def send_message(self,my_message="介绍一下自己吧"):
        #todo 发送消息
        # my_message = "我想写一篇关于使用Yolov10进行石头检测的论文，可以帮我推荐几篇参考文献嘛"
        # my_message=my_message+" 400字以内"
        try:
            e_input = self.br.execute_script(
                'return document.querySelector(\'p[data-placeholder="Message Claude..."]\')')

            e_input.send_keys(Keys.CONTROL, "a")
            e_input.send_keys(Keys.BACKSPACE)

            e_input.send_keys(my_message)
            e_input.send_keys(Keys.ENTER)

            time.sleep(2)
        except Exception as e:
            print(e)
        pass
    def get_respone(self):
        #todo 得到回复
        #获得消息
        try:
            e_replys = self.br.execute_script(
                'return document.querySelectorAll(\'div[class="contents"]\')')
            # print(len(e_replys))
            return self.restr(e_replys[-2].text)
        except Exception as e:
            print(e)
            return "这个问题无法回复，请换个问题"

        # print()
    def reboot(self):
        self.br.quit()
        self.__init__()

    def shot(self):
        self.br.quit()


    def respone_process(self):
        #todo 消息处理，防止表情
        pass

if __name__ == '__main__':
    #todo  retry可能不会显示出来
    #todo 如果回复有代码，无法正确显示回复

    myclaude=Myclaude()
    # myclaude.init()
    a="服务启动失败"
    try:
        myclaude.get_url("123号")
        myclaude.wait()
        # time.sleep(10)
        myclaude.send_message("讲个故事")
        myclaude.wait()
        # time.sleep(10)
        a=myclaude.get_respone()
        # time.sleep(10)
    except Exception as e:
        print(e)
        myclaude.shot()
    print(a)
    myclaude.shot()

